import os
import argparse
import random
import shutil
from shutil import copyfile
import cv2


def rm_mkdir(dir_path):
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
        print('Remove path - %s' % dir_path)
    os.makedirs(dir_path)
    print('Create path - %s' % dir_path)


def copy_file_difftype(src, dst, std_size):
    img = cv2.imread(src)
    img = cv2.resize(img, std_size)
    print(dst)
    cv2.imwrite(dst, img)


def main(config, std_size):
    train_path_ori = config.train_path
    valid_path_ori = config.valid_path
    test_path_ori = config.test_path
    origin_data_path_ori = config.origin_data_path
    sub_dir = os.listdir(origin_data_path_ori)

    for i in range(0, len(sub_dir)):
        train_path = os.path.join(train_path_ori, sub_dir[i])
        valid_path = os.path.join(valid_path_ori, sub_dir[i])
        test_path = os.path.join(test_path_ori, sub_dir[i])
        rm_mkdir(train_path)
        rm_mkdir(valid_path)
        rm_mkdir(test_path)
        origin_data_path = os.path.join(origin_data_path_ori, sub_dir[i])
        filenames = os.listdir(origin_data_path)
        data_list = []

        for filename in filenames:
            data_list.append(filename)

        num_total = len(data_list)
        num_train = int((config.train_ratio / (config.train_ratio + config.valid_ratio + config.test_ratio)) * num_total)
        num_valid = int((config.valid_ratio / (config.train_ratio + config.valid_ratio + config.test_ratio)) * num_total)
        num_test = num_total - num_train - num_valid

        print(origin_data_path)
        print('\nNum of train set : ', num_train)
        print('\nNum of valid set : ', num_valid)
        print('\nNum of test set : ', num_test)

        Arange = list(range(num_total))
        random.shuffle(Arange)

        for i in range(num_train):
            idx = Arange.pop()

            src = os.path.join(origin_data_path, data_list[idx])
            dst = os.path.join(train_path, data_list[idx])
            copy_file_difftype(src, dst, std_size)

        for i in range(num_valid):
            idx = Arange.pop()

            src = os.path.join(origin_data_path, data_list[idx])
            dst = os.path.join(valid_path, data_list[idx])
            copy_file_difftype(src, dst, std_size)

        for i in range(num_test):
            idx = Arange.pop()
            src = os.path.join(origin_data_path, data_list[idx])
            dst = os.path.join(test_path, data_list[idx])
            copy_file_difftype(src, dst, std_size)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train_ratio', type=float, default=0.8)
    parser.add_argument('--valid_ratio', type=float, default=0.1)
    parser.add_argument('--test_ratio', type=float, default=0.1)

    parser.add_argument('--origin_data_path', type=str,
                        default=r'E:\samples\dog_cat\train')
    # default = r'E:\xianhuang0818\3_cls_ori\enhance')
    parser.add_argument('--train_path', type=str,
                        default=r'E:\samples\dog_cat\for_train\train')
    parser.add_argument('--valid_path', type=str,
                        default=r'E:\samples\dog_cat\for_train\val')
    parser.add_argument('--test_path', type=str,
                        default=r'E:\samples\dog_cat\for_train\test')
    config = parser.parse_args()
    print(config)
    std_size = (224,224)
    main(config, std_size)